import numpy as np
import pickle
from tensorflow.keras.applications.vgg16 import VGG16, preprocess_input
from tensorflow.keras.preprocessing import image
import gradio as gr

# 加载预训练的 VGG16 模型（仅用于特征提取）
vgg_model = VGG16(weights='imagenet', include_top=False, pooling="max")

with open("best_Accuracy_model.pkl", 'rb') as f:
    svm_model = pickle.load(f)

def preprocess_image(image_path):
    # 加载并调整图像大小
    img = image.load_img(image_path, target_size=(224, 224))
    # 将图像转换为数组
    img_array = image.img_to_array(img)
    # 预处理图像以匹配 VGG16 的输入要求
    img_array = np.expand_dims(img_array, axis=0)
    img_array = preprocess_input(img_array)
    return img_array

def predict_image(image_path):
    # 预处理图像
    img_array = preprocess_image(image_path)
    # 使用 VGG16 提取特征
    features = vgg_model.predict(img_array)
    # 使用 SVM 模型进行预测
    prediction = svm_model.predict(features)
    # 将预测结果转换为猫或狗
    label = '狗' if prediction[0] == 1 else '猫'
    return label

# 创建 Gradio 接口
def main():
    iface = gr.Interface(
        fn=predict_image,
        inputs=gr.inputs.Image(type="filepath"),
        outputs=gr.outputs.Textbox()
    )
    iface.launch(debug=True)  # 设置为 False 以在生产环境中隐藏调试信息

if __name__ == "__main__":
    main()